import cv2
import onnxruntime as ort
import numpy as np
import sys
from itertools import product as product
from math import ceil
import torch

class PriorBox(object):
    def __init__(self,  image_size=None):
        super(PriorBox, self).__init__()
        self.min_sizes = [[10, 16, 24], [32, 48], [64, 96], [128, 192, 256]]
        self.steps = [8, 16, 32, 64]
        self.clip = False
        self.image_size = image_size
        self.feature_maps = [[ceil(self.image_size[0]/step), ceil(self.image_size[1]/step)] for step in self.steps]
        self.name = "s"

    def forward(self):
        anchors = []
        for k, f in enumerate(self.feature_maps):
            min_sizes = self.min_sizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / self.image_size[1]
                    s_ky = min_size / self.image_size[0]
                    dense_cx = [x * self.steps[k] / self.image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * self.steps[k] / self.image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]

        # back to torch land
        output = torch.Tensor(anchors).view(-1, 4)
        if self.clip:
            output.clamp_(max=1, min=0)
        return output

class FaceDetection:
    def __init__(self,file_path):
        self.ort_session = ort.InferenceSession(file_path)
        self.input_name  = self.ort_session.get_inputs()[0].name
        self.target_size = 640
    def _pre_processing(self, img):
        img_rgb = img[...,::-1] 
        img_rgb -= (104, 117, 123)
        img_rgb = img_rgb.transpose(2, 0, 1)
        return img_rgb[None, ...].astype(np.float32)    
    def run(self,img):
        img = img.astype(np.float32)
        input_tensor = self._pre_processing(img)
        loc, conf, landms = self.ort_session.run(None, {self.input_name: input_tensor}) 

        img_height = img.shape[0]
        img_width  = img.shape[1]
        priorbox = PriorBox( image_size=(img_height, img_width))
        priors = priorbox.forward()
        prior_data = priors.data

if __name__ == "__main__":
    model = FaceDetection(r'face_detect\weights\withlandmark\faceDetector.onnx')
    img = cv2.imread(r'face_detect\37.jpg')
    model.run(img)